import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载MNIST测试数据集
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 加载预训练模型（这里我们假设你有一个预训练的模型文件 'mnist_model.pth'）
model = torch.load('mnist_model.pth')
model.eval()  # 设置模型为评估模式

# 推理验证
correct = 0
total = 0
with torch.no_grad():  # 禁用梯度计算以加速推理过程
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f'Accuracy of the model on the MNIST test images: {accuracy}%')
